import org.bouncycastle.crypto.SecretWithEncapsulation;
import org.bouncycastle.pqc.crypto.mlkem.*;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.util.encoders.Hex;
import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;
import org.bouncycastle.crypto.digests.SHA256Digest;

public class Main {

    public static void main(String[] args) {
        try {
            // 测试ML-KEM-512，ML-KEM-768和ML-KEM-1024
            System.out.println("=== ML-KEM-512 加密解密示例 ===");
            testMLKEMEncryption(MLKEMParameters.ml_kem_512);

            System.out.println("=== ML-KEM-768 加密解密示例 ===");
            testMLKEMEncryption(MLKEMParameters.ml_kem_768);

            System.out.println("\n=== ML-KEM-1024 加密解密示例 ===");
            testMLKEMEncryption(MLKEMParameters.ml_kem_1024);

        } catch (Exception e) {
            System.err.println("发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    public static void testMLKEMEncryption(MLKEMParameters parameters) {
        SecureRandom random = new SecureRandom();

        try {
            // 1. 生成ML-KEM密钥对
            System.out.println("1. 生成ML-KEM密钥对...");
            MLKEMKeyPairGenerator keyPairGen = new MLKEMKeyPairGenerator();
            keyPairGen.init(new MLKEMKeyGenerationParameters(random, parameters));
            AsymmetricCipherKeyPair keyPair = keyPairGen.generateKeyPair();

            MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters) keyPair.getPublic();
            MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters) keyPair.getPrivate();
            System.out.println("公钥（字节:"+publicKey.getEncoded().length+"）:" + Hex.toHexString(publicKey.getEncoded()));
            System.out.println("私钥（字节:"+privateKey.getEncoded().length +"）:" + Hex.toHexString(privateKey.getEncoded()));
            // 2. 要加密的原始数据
            String originalData = "这是一条需要加密的敏感数据，包含重要信息！";
            System.out.println("原始数据: " + originalData);

            // 3. 加密过程
            System.out.println("\n2. 开始加密过程...");
            EncryptionResult encryptionResult = encryptData(publicKey, originalData.getBytes());

            System.out.println("封装后的密文长度: " + encryptionResult.encapsulatedKey.length + " 字节");
            System.out.println("数据密文长度: " + encryptionResult.encryptedData.length + " 字节");
            System.out.println("使用的共享密钥: " +
                    Hex.toHexString(encryptionResult.sharedSecret));

            // 4. 解密过程
            System.out.println("\n3. 开始解密过程...");
            byte[] decryptedData = decryptData(privateKey,
                    encryptionResult.encapsulatedKey,
                    encryptionResult.encryptedData,
                    encryptionResult.iv);

            String decryptedText = new String(decryptedData);
            System.out.println("解密后的数据: " + decryptedText);

            // 5. 验证结果
            boolean success = originalData.equals(decryptedText);
            System.out.println("\n4. 验证结果: " + (success ? "✓ 加解密成功！" : "✗ 加解密失败！"));

        } catch (Exception e) {
            System.err.println("ML-KEM加解密测试失败: " + e.getMessage());
            e.printStackTrace();
        }
    }
    public static byte[] hashTo32Bytes(byte[] encapsulatedKey) {
        SHA256Digest digest = new SHA256Digest();
        byte[] result = new byte[32];
        digest.update(encapsulatedKey, 0, encapsulatedKey.length);
        digest.doFinal(result, 0);
        return result;
    }
    /**
     * 加密数据
     * 使用ML-KEM封装共享密钥，然后用共享密钥进行AES-GCM加密
     */
    public static EncryptionResult encryptData(MLKEMPublicKeyParameters publicKey, byte[] data) throws Exception {
        SecureRandom random = new SecureRandom();

        // 1. 使用ML-KEM公钥封装共享密钥
        MLKEMGenerator kemGenerator = new MLKEMGenerator(random);
        SecretWithEncapsulation encapsulated = kemGenerator.generateEncapsulated(publicKey);

        byte[] encapsulatedKey = encapsulated.getEncapsulation();
        byte[] sharedSecret =  encapsulated.getSecret();

        // 2. 使用共享密钥进行AES-GCM加密
        // 创建AES密钥
        SecretKeySpec aesKey = new SecretKeySpec(sharedSecret, "AES");

        // 初始化加密器
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        byte[] iv = new byte[12]; // GCM推荐使用12字节的IV
        random.nextBytes(iv);
        GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv); // 128位认证标签
        cipher.init(Cipher.ENCRYPT_MODE, aesKey, parameterSpec);

        // 加密数据
        byte[] encryptedData = cipher.doFinal(data);

        return new EncryptionResult(encapsulatedKey, encryptedData, sharedSecret, iv);
    }

    /**
     * 解密数据
     * 使用ML-KEM私钥解封装共享密钥，然后用共享密钥进行AES-GCM解密
     */
    public static byte[] decryptData(MLKEMPrivateKeyParameters privateKey,
                                     byte[] encapsulatedKey,
                                     byte[] encryptedData,
                                     byte[] iv) throws Exception {

        // 1. 使用ML-KEM私钥解封装共享密钥
        MLKEMExtractor kemExtractor = new MLKEMExtractor(privateKey);
        byte[] sharedSecret = kemExtractor.extractSecret(encapsulatedKey);

        // 2. 使用共享密钥进行AES-GCM解密
        SecretKeySpec aesKey = new SecretKeySpec(sharedSecret, "AES");

        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv);
        cipher.init(Cipher.DECRYPT_MODE, aesKey, parameterSpec);

        return cipher.doFinal(encryptedData);
    }

    /**
     * 加密结果容器类
     */
    public static class EncryptionResult {
        public final byte[] encapsulatedKey;  // ML-KEM封装后的密钥
        public final byte[] encryptedData;    // AES加密后的数据
        public final byte[] sharedSecret;     // 共享密钥
        public final byte[] iv;               // 初始化向量

        public EncryptionResult(byte[] encapsulatedKey, byte[] encryptedData,
                                byte[] sharedSecret, byte[] iv) {
            this.encapsulatedKey = encapsulatedKey;
            this.encryptedData = encryptedData;
            this.sharedSecret = sharedSecret;
            this.iv = iv;
        }
    }

    /**
     * 保存和加载密钥的示例方法
     */
    public static void keyManagementExample() {
        try {
            SecureRandom random = new SecureRandom();

            // 生成密钥对
            MLKEMKeyPairGenerator keyPairGen = new MLKEMKeyPairGenerator();
            keyPairGen.init(new MLKEMKeyGenerationParameters(random, MLKEMParameters.ml_kem_768));
            AsymmetricCipherKeyPair keyPair = keyPairGen.generateKeyPair();

            MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters) keyPair.getPublic();
            MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters) keyPair.getPrivate();

            // 将密钥序列化为字节数组（用于存储或传输）
            byte[] publicKeyBytes = publicKey.getEncoded();
            byte[] privateKeyBytes = privateKey.getEncoded();

            System.out.println("公钥Hex: " + Hex.toHexString(publicKeyBytes));
            System.out.println("私钥Hex: " + Hex.toHexString(privateKeyBytes));

            // 注意：在实际应用中，私钥应该安全存储（如使用密钥库或加密存储）

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}